package middleware

import (
	"bytes"
	"fmt"
	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/core/stores/sqlx"
	"io"
	"net"
	"net/http"
	"time"
	"yunzhan/api-server/internal/config"
	"yunzhan/common/utils"
)

type AuditMiddleware struct {
	Conn sqlx.SqlConn
	Conf config.Config
}

type responseWriter struct {
	http.ResponseWriter
	body       *bytes.Buffer
	statusCode int
}

func (rw *responseWriter) Write(b []byte) (int, error) {
	rw.body.Write(b)
	return rw.ResponseWriter.Write(b)
}

func (rw *responseWriter) WriteHeader(statusCode int) {
	rw.statusCode = statusCode
	rw.ResponseWriter.WriteHeader(statusCode)
}

func NewAuditMiddleware(conn sqlx.SqlConn, c config.Config) *AuditMiddleware {
	return &AuditMiddleware{Conn: conn, Conf: c}
}

func (m *AuditMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		clientIP, _, err := net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			logx.Errorf("Error extracting IP: %v", err)
			clientIP = "unknown"
		}

		operationType := r.Method
		requestURI := r.RequestURI
		userAgent := r.UserAgent()

		serverNode, err := utils.GetCurrentInternalIP()
		if err != nil {
			logx.Errorf("获取主机ip失败: %v", err)
			serverNode = "unknown"
		}

		requestBody, err := io.ReadAll(r.Body)
		if err != nil {
			logx.Errorf("Error reading request body: %v", err)
		}
		// 重置请求体
		r.Body = io.NopCloser(bytes.NewBuffer(requestBody))
		rw := &responseWriter{
			ResponseWriter: w,
			body:           new(bytes.Buffer),
			statusCode:     http.StatusOK,
		}

		// 使用包装的响应写入器
		next(rw, r)

		end := time.Now()
		duration := end.Sub(start)

		statusCode := rw.statusCode
		responseBody := rw.body.String()

		logx.Infof("Audit Log - ClientIP: %s, Operation: %s, URI: %s, Start: %s, End: %s, Duration: %s, Status: %d",
			clientIP,
			operationType,
			requestURI,
			start.Format("2006-01-02 15:04:05.000"),
			end.Format("2006-01-02 15:04:05.000"),
			duration,
			statusCode,
		)

		tableName := m.Conf.SQLite.Tables.AuditLogs
		if m.Conf.UseMysql {
			tableName = m.Conf.Mysql.Tables.AuditLogs
		}

		_, err = m.Conn.Exec(
			fmt.Sprintf(`INSERT INTO %s 
             (client_ip, operation_type, request_uri, start_time, end_time, duration, status_code, request_body, response_body, user_agent, server_node) 
             VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, tableName),
			clientIP, operationType, requestURI, start.Format("2006-01-02 15:04:05.000"), end.Format("2006-01-02 15:04:05.000"), duration.Milliseconds(), statusCode, requestBody, responseBody, userAgent, serverNode,
		)
		if err != nil {
			logx.Errorf("Failed to write audit log to database: %v", err)
		}
	}
}
